import pygame
from pygame.locals import *
from OpenGL.GL import *
from OpenGL.GLU import *
import numpy as np
import random
import math
import sys
import time

# 添加调试信息
def print_debug(message):
    print(f"[DEBUG] {message}")

# 初始化Pygame
pygame.init()
WIDTH, HEIGHT = 1280, 720
print_debug("初始化Pygame完成")

try:
    display = pygame.display.set_mode((WIDTH, HEIGHT), DOUBLEBUF | OPENGL)
    pygame.display.set_caption("增强版3D网络 (NVIDIA GTX 1050)")
    print_debug("创建OpenGL窗口成功")
except Exception as e:
    print(f"OpenGL初始化错误: {e}")
    sys.exit(1)

# 检查OpenGL版本
opengl_version = glGetString(GL_VERSION).decode()
print(f"OpenGL版本: {opengl_version}")
print(f"OpenGL供应商: {glGetString(GL_VENDOR).decode()}")
print(f"OpenGL渲染器: {glGetString(GL_RENDERER).decode()}")

# 基本OpenGL设置
glEnable(GL_DEPTH_TEST)
glEnable(GL_BLEND)
glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)

# 设置光照
glEnable(GL_LIGHTING)
glEnable(GL_LIGHT0)
glEnable(GL_COLOR_MATERIAL)
glColorMaterial(GL_FRONT_AND_BACK, GL_AMBIENT_AND_DIFFUSE)

# 设置光照参数
light_position = [0.0, 0.0, 200.0, 1.0]
light_ambient = [0.3, 0.3, 0.3, 1.0]
light_diffuse = [0.8, 0.8, 0.8, 1.0]
light_specular = [1.0, 1.0, 1.0, 1.0]

glLightfv(GL_LIGHT0, GL_POSITION, light_position)
glLightfv(GL_LIGHT0, GL_AMBIENT, light_ambient)
glLightfv(GL_LIGHT0, GL_DIFFUSE, light_diffuse)
glLightfv(GL_LIGHT0, GL_SPECULAR, light_specular)

# 设置视角
glMatrixMode(GL_PROJECTION)
glLoadIdentity()
gluPerspective(45, (WIDTH/HEIGHT), 0.1, 1000.0)
glMatrixMode(GL_MODELVIEW)
glLoadIdentity()
glTranslatef(0.0, 0.0, -600)

print_debug("OpenGL初始设置完成")

# 使用顶点数组绘制球体
def create_sphere_vbo(radius, slices, stacks):
    vertices = []
    normals = []
    indices = []
    
    # 生成顶点和法线
    for i in range(stacks + 1):
        phi = math.pi * i / stacks
        for j in range(slices + 1):
            theta = 2.0 * math.pi * j / slices
            
            # 球体坐标
            x = radius * math.sin(phi) * math.cos(theta)
            y = radius * math.sin(phi) * math.sin(theta)
            z = radius * math.cos(phi)
            
            # 单位法线向量
            nx = x / radius
            ny = y / radius
            nz = z / radius
            
            vertices.extend([x, y, z])
            normals.extend([nx, ny, nz])
    
    # 生成索引
    for i in range(stacks):
        for j in range(slices):
            # 每个四边形的两个三角形
            first = i * (slices + 1) + j
            second = first + slices + 1
            
            # 第一个三角形
            indices.extend([first, second, first + 1])
            
            # 第二个三角形
            indices.extend([second, second + 1, first + 1])
    
    # 创建顶点缓冲对象
    vbo_vertices = glGenBuffers(1)
    glBindBuffer(GL_ARRAY_BUFFER, vbo_vertices)
    glBufferData(GL_ARRAY_BUFFER, 
                np.array(vertices, dtype=np.float32),
                GL_STATIC_DRAW)
    
    # 创建法线缓冲对象
    vbo_normals = glGenBuffers(1)
    glBindBuffer(GL_ARRAY_BUFFER, vbo_normals)
    glBufferData(GL_ARRAY_BUFFER, 
                np.array(normals, dtype=np.float32),
                GL_STATIC_DRAW)
    
    # 创建索引缓冲对象
    vbo_indices = glGenBuffers(1)
    glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, vbo_indices)
    glBufferData(GL_ELEMENT_ARRAY_BUFFER, 
                np.array(indices, dtype=np.uint32),
                GL_STATIC_DRAW)
    
    return {
        'vertices': vbo_vertices,
        'normals': vbo_normals,
        'indices': vbo_indices,
        'count': len(indices)
    }

# 球体类
class Ball:
    def __init__(self, index, sphere_data):
        self.index = index
        self.radius = random.uniform(3, 8)
        self.scale_factor = self.radius / 10.0  # 假设球体模型半径为10
        
        # 生成明亮颜色 - 根据索引分组，同组颜色相似
        group = index % 6
        base_h = group / 6.0
        h = base_h + random.uniform(-0.05, 0.05)  # 小范围随机变化
        s = random.uniform(0.7, 1)
        v = random.uniform(0.8, 1)
        r, g, b = self._hsv_to_rgb(h, s, v)
        
        # 设置颜色和透明度
        self.color = (r, g, b, 0.9)
        
        # 在球体内随机生成位置 - 使用改进的分布
        # 避免所有球都聚集在中心
        r = 170 * math.pow(random.random(), 0.5)  # 非均匀分布，靠近边缘更多
        theta = random.uniform(0, math.pi * 2)
        phi = random.uniform(0, math.pi)
        
        self.x = r * math.sin(phi) * math.cos(theta)
        self.y = r * math.sin(phi) * math.sin(theta)
        self.z = r * math.cos(phi)
        
        # 随机速度 - 稍微加快
        self.speed = [random.uniform(-1.2, 1.2) for _ in range(3)]
        
        # 存储球体VBO数据的引用
        self.sphere_data = sphere_data
        
        # 上次更新时间
        self.last_update = time.time()
    
    def _hsv_to_rgb(self, h, s, v):
        if s == 0.0:
            return v, v, v
        
        i = int(h * 6)
        f = (h * 6) - i
        p = v * (1 - s)
        q = v * (1 - s * f)
        t = v * (1 - s * (1 - f))
        i %= 6
        
        if i == 0: return v, t, p
        if i == 1: return q, v, p
        if i == 2: return p, v, t
        if i == 3: return p, q, v
        if i == 4: return t, p, v
        if i == 5: return v, p, q
    
    def update(self, dt):
        # 使用帧时间差来平滑动画
        dt = min(dt, 0.05)  # 防止时间间隔过大
        
        # 更新位置
        self.x += self.speed[0] * dt * 30  # 调整为每秒的速度
        self.y += self.speed[1] * dt * 30
        self.z += self.speed[2] * dt * 30
        
        # 碰撞检测
        container_radius = 190
        distance = math.sqrt(self.x**2 + self.y**2 + self.z**2)
        
        if distance + self.radius > container_radius:
            # 法线向量
            nx = self.x / distance
            ny = self.y / distance
            nz = self.z / distance
            
            # 速度点积
            dot = self.speed[0] * nx + self.speed[1] * ny + self.speed[2] * nz
            
            # 反弹 - 添加一点能量损失
            loss = 0.98  # 2%能量损失
            self.speed[0] = (self.speed[0] - 2 * dot * nx) * loss
            self.speed[1] = (self.speed[1] - 2 * dot * ny) * loss
            self.speed[2] = (self.speed[2] - 2 * dot * nz) * loss
            
            # 调整位置
            adj_radius = container_radius - self.radius - 1
            self.x = nx * adj_radius
            self.y = ny * adj_radius
            self.z = nz * adj_radius
            
        # 更新时间
        self.last_update = time.time()
    
    def draw(self):
        glPushMatrix()
        glTranslatef(self.x, self.y, self.z)
        glScalef(self.scale_factor, self.scale_factor, self.scale_factor)
        
        # 设置颜色
        glColor4f(*self.color)
        
        # 启用顶点和法线数组
        glEnableClientState(GL_VERTEX_ARRAY)
        glEnableClientState(GL_NORMAL_ARRAY)
        
        # 绑定顶点缓冲对象
        glBindBuffer(GL_ARRAY_BUFFER, self.sphere_data['vertices'])
        glVertexPointer(3, GL_FLOAT, 0, None)
        
        # 绑定法线缓冲对象
        glBindBuffer(GL_ARRAY_BUFFER, self.sphere_data['normals'])
        glNormalPointer(GL_FLOAT, 0, None)
        
        # 绑定索引缓冲对象并绘制
        glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, self.sphere_data['indices'])
        glDrawElements(GL_TRIANGLES, self.sphere_data['count'], GL_UNSIGNED_INT, None)
        
        # 禁用顶点和法线数组
        glDisableClientState(GL_VERTEX_ARRAY)
        glDisableClientState(GL_NORMAL_ARRAY)
        
        glPopMatrix()

def calculate_distance(ball1, ball2):
    return math.sqrt(
        (ball1.x - ball2.x)**2 + 
        (ball1.y - ball2.y)**2 + 
        (ball1.z - ball2.z)**2
    )

# 创建球体VBO
print_debug("创建球体VBO...")
try:
    # 使用较低精度的球体模型来提高性能
    sphere_data = create_sphere_vbo(10.0, 16, 16)
    print_debug("球体VBO创建成功")
except Exception as e:
    print(f"创建球体VBO失败: {e}")
    import traceback
    traceback.print_exc()
    pygame.quit()
    sys.exit(1)

# 创建小球
print("创建小球...")
num_balls = 200  # 增加到200个小球
try:
    balls = [Ball(i, sphere_data) for i in range(num_balls)]
    print(f"已创建 {num_balls} 个小球")
except Exception as e:
    print(f"创建小球失败: {e}")
    import traceback
    traceback.print_exc()
    pygame.quit()
    sys.exit(1)

# 旋转状态
rotation_x = 0
rotation_y = 0
rotation_scale = 1.0

# 保存好看的视角
favorite_views = [
    {'x': 15, 'y': 45, 'scale': 1.2},  # 偏上视角
    {'x': -20, 'y': 120, 'scale': 0.9},  # 偏下视角
    {'x': 0, 'y': 180, 'scale': 1.0},  # 正面视角
]
current_view = 0

# 主循环
running = True
clock = pygame.time.Clock()
pulse_time = 0
last_time = time.time()
auto_rotate = True

# 连接线渲染设置
connection_distance = 60  # 连接距离
max_connections = 10000  # 最大连接数

# 颜色周期效果
color_cycle_speed = 0.2
color_cycle_offset = 0

print("开始渲染...")
try:
    while running:
        # 计算帧时间差
        current_time = time.time()
        dt = current_time - last_time
        last_time = current_time
        
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    running = False
                elif event.key == pygame.K_SPACE:
                    # 切换自动旋转
                    auto_rotate = not auto_rotate
                elif event.key == pygame.K_v:
                    # 切换预设视角
                    current_view = (current_view + 1) % len(favorite_views)
                    view = favorite_views[current_view]
                    rotation_x = view['x']
                    rotation_y = view['y']
                    rotation_scale = view['scale']
            elif event.type == pygame.MOUSEBUTTONDOWN:
                if event.button == 4:  # 滚轮上滚
                    rotation_scale *= 1.1
                elif event.button == 5:  # 滚轮下滚
                    rotation_scale /= 1.1
                elif event.button == 1:  # 左键点击
                    auto_rotate = False
        
        # 鼠标控制旋转
        if pygame.mouse.get_pressed()[0]:  # 左键按下
            rel_x, rel_y = pygame.mouse.get_rel()
            rotation_y += rel_x * 0.5
            rotation_x += rel_y * 0.5
        else:
            pygame.mouse.get_rel()  # 重置相对位置
            # 自动旋转
            if auto_rotate:
                rotation_y += 0.3 * dt * 60  # 每秒15度
        
        # 清除屏幕
        glClearColor(0.0, 0.0, 0.05, 1.0)  # 深蓝色背景
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
        
        # 设置视角
        glLoadIdentity()
        glTranslatef(0.0, 0.0, -600 * rotation_scale)
        
        # 应用旋转
        glRotatef(rotation_x, 1, 0, 0)
        glRotatef(rotation_y, 0, 1, 0)
        
        # 更新小球
        for ball in balls:
            ball.update(dt)
        
        # 绘制容器
        glDisable(GL_LIGHTING)
        
        # 周期性变化容器颜色
        color_cycle_offset += color_cycle_speed * dt
        r = 0.4 + 0.1 * math.sin(color_cycle_offset)
        g = 0.4 + 0.1 * math.sin(color_cycle_offset + 2)
        b = 0.6 + 0.2 * math.sin(color_cycle_offset + 4)
        glColor4f(r, g, b, 0.3)
        
        # 使用wireframe方式绘制容器球体
        sphere = gluNewQuadric()
        gluQuadricDrawStyle(sphere, GLU_LINE)
        gluSphere(sphere, 200, 32, 32)
        gluDeleteQuadric(sphere)
        
        # 绘制连接线
        glLineWidth(1.5)
        
        glBegin(GL_LINES)
        line_count = 0
        for i, ball1 in enumerate(balls):
            # 限制每个球的连接数量，避免过度密集
            ball_connections = 0
            max_ball_connections = 7  # 每个球最多连接7个其他球
            
            for j, ball2 in enumerate(balls[i+1:], i+1):
                if line_count >= max_connections or ball_connections >= max_ball_connections:
                    break
                    
                distance = calculate_distance(ball1, ball2)
                if distance < connection_distance:
                    # 混合两个球的颜色
                    r = (ball1.color[0] + ball2.color[0]) / 2
                    g = (ball1.color[1] + ball2.color[1]) / 2
                    b = (ball1.color[2] + ball2.color[2]) / 2
                    
                    # 距离越远，线越透明
                    alpha = 0.8 * (1.0 - distance / connection_distance)
                    
                    # 距离越远，线越细
                    glColor4f(r, g, b, alpha)
                    glVertex3f(ball1.x, ball1.y, ball1.z)
                    glVertex3f(ball2.x, ball2.y, ball2.z)
                    line_count += 1
                    ball_connections += 1
        glEnd()
        
        # 重新启用光照
        glEnable(GL_LIGHTING)
        
        # 脉动效果
        pulse_time += 0.8 * dt
        ambient = 0.2 + 0.1 * math.sin(pulse_time)
        light_ambient = [ambient, ambient, ambient, 1.0]
        glLightfv(GL_LIGHT0, GL_AMBIENT, light_ambient)
        
        # 移动光源
        light_x = 200 * math.sin(pulse_time * 0.3)
        light_y = 100 * math.sin(pulse_time * 0.5)
        light_z = 200 * math.cos(pulse_time * 0.3)
        light_position = [light_x, light_y, light_z, 1.0]
        glLightfv(GL_LIGHT0, GL_POSITION, light_position)
        
        # 绘制所有小球
        for ball in balls:
            ball.draw()
        
        # 更新显示
        pygame.display.flip()
        
        # 显示帧率和线条数量
        fps = clock.get_fps()
        pygame.display.set_caption(f"增强版3D网络 - FPS: {fps:.1f} - 线条: {line_count} - 小球: {num_balls}")
        
        # 控制帧率
        clock.tick(60)

except Exception as e:
    print(f"渲染循环中发生错误: {e}")
    import traceback
    traceback.print_exc()

# 清理资源
glDeleteBuffers(3, [sphere_data['vertices'], sphere_data['normals'], sphere_data['indices']])

# 退出
pygame.quit()
print("程序结束") 